package flink.dataset;

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.junit.Test;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author binarylei
 * @version 2019-11-29
 */
public class BroadcastDataSetTest {

    /**
     * 文件变量，如果多个算子需要使用同一个变量，需要在算子后注册广播变量
     * <pre>
     * eg: dataset.map().withBroadcastSet(agesToBroadcast, "agesToBroadcast")
     *            .map().withBroadcastSet(agesToBroadcast, "agesToBroadcast")
     * </pre>
     * 注意：
     * 1. 广播出去的变量存在于每个节点的内存中，所以这个数据集不能太大。因为会常驻内存，除非程序执行结束。
     * 2. 广播变量在初始化广播出去后不支持修改，这样才能保证每个节点数据都是一致的。
     *
     * @throws Exception
     */
    @Test
    public void testBroadcast() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        // 准备需要广播的数据
        // <id, age>
        DataSet<Tuple2<Long, Integer>> agesToBroadcast = env.fromElements(
                new Tuple2<Long, Integer>(1L, 20),
                new Tuple2<Long, Integer>(2L, 20),
                new Tuple2<Long, Integer>(4L, 40));

        // <id, name>
        DataSource<Tuple2<Long, String>> names = env.fromElements(
                new Tuple2<Long, String>(1L, "john"),
                new Tuple2<Long, String>(2L, "mick"),
                new Tuple2<Long, String>(3L, "bai"));
        names.map(new RichMapFunction<Tuple2<Long, String>, Tuple3<Long, String, Integer>>() {
            Map<Long, Integer> broadcastMap = new HashMap<>();

            /**
             * 方法只会执行一次，可以执行初始化功能
             */
            @Override
            public void open(Configuration parameters) throws Exception {
                super.open(parameters);
                // 获取广播变量
                List<Tuple2<Long, Integer>> broadcastList = getRuntimeContext()
                        .getBroadcastVariable("agesToBroadcast");
                broadcastList.forEach(t -> broadcastMap.put(t.f0, t.f1));
            }

            @Override
            public Tuple3<Long, String, Integer> map(Tuple2<Long, String> t)
                    throws Exception {
                Integer age = broadcastMap.get(t.f0);
                return new Tuple3<>(t.f0, t.f1, age != null ? age : -1);
            }
        }).withBroadcastSet(agesToBroadcast, "agesToBroadcast").print();

    }
}
